import cv2
import numpy as np
from ultralytics import YOLO
import time


class VehicleDetectionModule:
    def __init__(self, warning_threshold=30):
        # 加载YOLOv8模型
        self.model = YOLO("yolov8n.pt")
        # 车辆类别ID (YOLOv8中car=2, truck=7, bus=5)
        self.vehicle_classes = [2, 5, 7]
        self.vehicle_names = {2: "轿车", 5: "公交车", 7: "卡车"}

        # 预警阈值(米)
        self.warning_threshold = warning_threshold

        # 相机内参(示例值，实际需校准)
        self.focal_length = 800  # 焦距
        self.vehicle_real_height = 1.5  # 车辆平均高度(米)

        # 状态变量
        self.vehicle_count = 0
        self.closest_distance = 9999.0
        self.warning = False

    def estimate_distance(self, vehicle_height_pixels):
        """根据像素高度估算距离"""
        if vehicle_height_pixels <= 0:
            return 9999.0
        # 距离 = (实际高度 * 焦距) / 像素高度
        return (self.vehicle_real_height * self.focal_length) / vehicle_height_pixels

    def process_frame(self, frame):
        """处理单帧图像，返回检测结果"""
        # 重置状态
        self.vehicle_count = 0
        self.closest_distance = 9999.0
        self.warning = False

        # 目标检测
        results = self.model(frame, classes=self.vehicle_classes)

        # 处理检测结果
        annotated_frame = frame.copy()
        for result in results:
            for box in result.boxes:
                self.vehicle_count += 1

                # 获取边界框
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                vehicle_height = y2 - y1  # 车辆像素高度

                # 估算距离
                distance = self.estimate_distance(vehicle_height)
                self.closest_distance = min(self.closest_distance, distance)

                # 绘制边界框和信息
                class_id = int(box.cls[0])
                vehicle_type = self.vehicle_names.get(class_id, "车辆")
                confidence = float(box.conf[0])

                label = f"{vehicle_type} {confidence:.2f} 距离: {distance:.1f}m"
                color = (0, 0, 255) if distance < self.warning_threshold else (0, 255, 0)

                cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
                cv2.putText(annotated_frame, label, (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        # 判断是否需要预警
        if self.closest_distance < self.warning_threshold and self.vehicle_count > 0:
            self.warning = True
            cv2.putText(annotated_frame, "⚠️ 前方车辆过近，请谨慎驾驶！",
                        (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 3)

        # 返回处理结果
        return {
            "frame": annotated_frame,
            "vehicle_count": self.vehicle_count,
            "closest_distance": self.closest_distance,
            "warning": self.warning
        }


# 测试代码
if __name__ == "__main__":
    # 初始化模块
    vehicle_module = VehicleDetectionModule(warning_threshold=30)

    # 打开摄像头
    cap = cv2.VideoCapture(0)  # 0为默认摄像头，可替换为视频文件路径

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # 处理帧
        result = vehicle_module.process_frame(frame)

        # 显示结果
        cv2.imshow("车辆检测与距离预警", result["frame"])

        # 按q退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
